package com.zhipin.hackthon;

import java.util.Arrays;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

public class ZslMultiAtomicInteger implements MultiAtomicInteger {

    /**
     * lock集合
     */
    private final Object[] locks;

    /**
     * [hash]
     */
    private final long[][] map;

    /**
     * 0有特殊含义，单独用一个变量记
     */
    private final AtomicInteger zeroNum = new AtomicInteger();
    /**
     * 标记0是否放入过
     */
    private final AtomicBoolean existsZero = new AtomicBoolean();

    private final AtomicInteger size = new AtomicInteger();

    public ZslMultiAtomicInteger(int capacity) {
        // 系数越大，吞吐量越高，内存也越多
//        capacity = (int) Math.max(1, capacity * 0.35);
        capacity = (int) Math.max(1, capacity * 0.4);

        // 理论上锁数量越多，吞吐量越高，但是太多的话，占用内存也多
        int lockNum = Math.min(128, capacity);
        locks = new Object[lockNum];
        for (int i = 0; i < locks.length; i++) {
            locks[i] = new Object();
        }

        map = new long[capacity][];
    }

    @Override
    public int addAndGet(int key, int cnt) {
        if (key == 0) {
            cnt = zeroNum.addAndGet(cnt);
            boolean exists = existsZero.getAndSet(true);
            if (!exists) {
                size.incrementAndGet();
            }
            return cnt;
        }

        int hash = hash(key);
        int lockIndex = hash % locks.length;
        synchronized (locks[lockIndex]) {
            long[] arr = map[hash];
            if (arr == null) {
                arr = map[hash] = new long[2]; // 初始化链表大小2
            }

            for (int i = 0; i < arr.length; i++) {
                if (arr[i] == 0) {
                    arr[i] = createNode(key, cnt);
                    size.incrementAndGet();
                    return cnt;
                }

                if (key == getKey(arr[i])) {
                    arr[i] += cnt;
                    return getValue(arr[i]);
                }
            }

            int oldLen = arr.length;
            arr = map[hash] = Arrays.copyOf(arr, Math.min(arr.length << 1, Integer.MAX_VALUE));
            arr[oldLen] = createNode(key, cnt);
            size.incrementAndGet();
            return cnt;
        }
    }

    @Override
    public int get(int key) {
        if (key == 0) {
            return zeroNum.get();
        }

        int hash = hash(key);
        long[] arr = map[hash];
        if (arr == null) {
            return 0;
        }

        for (long node : arr) {
            if (node == 0) {
                return 0;
            }

            if (key == getKey(node)) {
                return getValue(node);
            }
        }

        return 0;
    }

    @Override
    public int size() {
        return size.get();
    }

    private int hash(int key) {
        return key % map.length;
    }

    private long createNode(int key, int value) {
        return (long) key << 32 | value;
    }

    private int getKey(long node) {
        return (int) (node >> 32);
    }

    private int getValue(long node) {
        return (int) (node);
    }
}
